
#include "pows_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
#include <algorithm>

//UB:253952
namespace optiling {
    const uint32_t BLOCK_SIZE = 512;
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

  PowsTilingData tiling;
    const gert::StorageShape* x1_StorageShape = context->GetInputShape(0);
    const gert::StorageShape* x2_StorageShape = context->GetInputShape(1);

    uint32_t x1_size = x1_StorageShape->GetStorageShape().GetShapeSize();
    uint32_t x2_size = x2_StorageShape->GetStorageShape().GetShapeSize();

    auto dt = context->GetInputTensor(1)->GetDataType();
    
    if(x1_size != x2_size) //广播
    {
        uint32_t x1_dim = x1_StorageShape->GetStorageShape().GetDimNum();
        uint32_t x2_dim = x2_StorageShape->GetStorageShape().GetDimNum();
        uint32_t y_size;
        uint32_t y_dim;
        uint32_t x1_shape[8];
        uint32_t x2_shape[8];
        uint32_t y_shape[8];
        y_dim = x1_dim;
        if(y_dim < x2_dim) y_dim = x2_dim;

        y_size = x1_size;
        if(y_size < x2_size) y_size = x2_size;

        uint32_t x1_start = y_dim - x1_dim;
        uint32_t x2_start = y_dim - x2_dim;

        for(int i=0; i< y_dim; i++)
        {
            if(i < x1_start) x1_shape[i] = 1;
            else x1_shape[i] = x1_StorageShape->GetStorageShape().GetDim(i - x1_start);

            if(i < x2_start) x2_shape[i] = 1;
            else x2_shape[i] = x2_StorageShape->GetStorageShape().GetDim(i - x2_start);

            y_shape[i] = x1_shape[i];
            if(y_shape[i] < x2_shape[i]) y_shape[i] = x2_shape[i];
        }
        tiling.set_size(y_size);
        tiling.set_y_dim(y_dim);
        tiling.set_x1_shape(x1_shape);
        tiling.set_x2_shape(x2_shape);
        tiling.set_y_shape(y_shape);

        context->SetTilingKey(2);
    }
    else
    {
        uint64_t ubSize;
        auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
        ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize); //获取硬件平台存储空间 UB 的内存大小
        
        uint32_t inputNum = context->GetInputShape(1)->GetStorageShape().GetShapeSize(); //输入数量
        uint32_t inputBytes = GetSizeByDataType(context->GetInputDesc(0)->GetDataType()); //输入类型
        uint32_t inputLength = inputBytes * inputNum; //输入长度

        // uint32_t ubDataNumber = 3*2+4;
        // if (dt == ge::DT_FLOAT) ubDataNumber = 3*2;

        // ubDataNumber *= 2;


        uint32_t ubDataNumber = 3*2*2+2*4;
        if (dt == ge::DT_FLOAT) ubDataNumber = 3*2*4;

        uint32_t alignment = BLOCK_SIZE / inputBytes;
        uint32_t alignment_32 = 32 / inputBytes;
        uint32_t tileDataNum = 10496;
        uint32_t TailDataNum = (((inputNum % tileDataNum) + alignment_32 - 1) / alignment_32)* alignment_32;
        uint32_t finalTileNum = (TailDataNum == 0) ? (inputNum / tileDataNum) : (inputNum / tileDataNum + 1);
        uint32_t CoreDataNum = ((inputNum + alignment - 1) / alignment)* alignment;

        TailDataNum = (TailDataNum == 0) ? (tileDataNum) : (TailDataNum);

        tiling.set_CoreDataNum(CoreDataNum);  //对齐空间后的输入数量
        tiling.set_finalTileNum(finalTileNum-1);//需要循环处理几次
        tiling.set_tileDataNum(tileDataNum); //每次处理的数据量
        tiling.set_TailDataNum(TailDataNum); //最后一次需要处理的数据量

        context->SetTilingKey(1);
    }
  context->SetBlockDim(1);
  tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
  context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

  return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class Pows : public OpDef {
public:
    explicit Pows(const char* name) : OpDef(name)
    {
        this->Input("x1")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("x2")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");

    }
};

OP_ADD(Pows);
}
